package net.tp.algo.trie;

import java.util.ArrayList;
import java.util.List;
import java.util.Stack;

/**
 * 
 * Ternary Search Trie.
 * 
 * Ref: Fast Algorithms for Sorting and Searching Strings
 * 
 * @author Trung Phan
 *
 */
public class TernaryTrie<V> implements Trie<V> {

	
	private Node<V> root;
	
	private int size;
	
	
	public int size() {
		return size;
	}
	
	public boolean empty() {
		return size == 0;
	}
	
	public void put(String key, V value) {
		if (value == null || key == null || key.length() == 0) {
			return;
		}
		
		
		if (root == null) {
			root = new Node<>(key.charAt(0), null);
		}
		
		Node<V> n = root;
		int i = 0;
		while (true) {
			char c = key.charAt(i);
			if (c < n.c) {
				if (n.left == null) {
					n.left = new Node<>(c, n);
				}
				n = n.left;
			}
			else if (c > n.c) {
				if (n.right == null) {
					n.right = new Node<>(c, n);
				}
				n = n.right;
			}
			else if (i == key.length() - 1) {
				break;
			}
			else {
				if (n.mid == null) {
					n.mid = new Node<>(c, n);
				}
				n = n.mid;
				i++;
			}
		}

		if (n.value == null) {
			size++;
		}
		n.value = value;
	}
	
	public V get(String key) {
		if (key == null || key.length() == 0) {
			return null;
		}
		
		Node<V> n = findNode(key);

		return n == null ? null : n.value;
	}

	private Node<V> findNode(String key) {
		Node<V> n = root;
		int i = 0;
		while (n != null) {
			char c = key.charAt(i);
			if (c < n.c) {
				n = n.left;
			}
			else if (c > n.c) {
				n = n.right;
			}
			else if (i == key.length() - 1) {
				break;
			}
			else {
				n = n.mid;
				i++;
			}
		}
		return n;
	}
	
	public V delete(String key) {
		if (key == null || key.length() == 0) {
			return null;
		}
		
		Node<V> n = findNode(key);

		if (n == null) {
			return null;
		}
		
		V result = n.value;
		if (n.value != null) {
			size--;
		}
		n.value = null;
		
		if (size == 0) {
			root = null;
		}
		else {
			while (n != null && n.isEmpty()) {
				n.parent.removeChild(n);
				n = n.parent;
			}
		}

		return result;
	}
	
	@Override
	public Iterable<String> keys() {
		final List<String> result = new ArrayList<>();
		dfs("", root, new Visitor<V>() {
			public void visit(String key, V value) {
				result.add(key);
			}
		});
		return result;
	}

	@Override
	public Iterable<String> keysWithPrefix(String s) {
		
		final List<String> result = new ArrayList<>();
		
		Node<V> n = findNode(s);
		if (n != null) {
			if (n.value != null) {
				result.add(s);
			}
			
			dfs(s, n.mid, new Visitor<V>() {
				public void visit(String key, V value) {
					result.add(key);
				}
			});
		}

		return result;
	}

	@Override
	public String longestPrefixOf(String s) {
		
		String result = "";
		
		Node<V> n = root;
		int i = 0;
		while (n != null) {
			char c = s.charAt(i);
			if (c < n.c) {
				n = n.left;
			}
			else if (c > n.c) {
				n = n.right;
			}
			else {
				if (n.value != null) {
					result = s.substring(0, i+1);
				}
				
				n = n.mid;
				i++;
			}
		}

		return result;
	}

	private static <V> void dfs(String keyPrefix, Node<V> root, Visitor<V> visitor) {
		if (root == null) {
			return;
		}
		Stack<StackFrame<V>> stack = new Stack<>();
		stack.push(new StackFrame<V>(root));
		StringBuilder keyBuilder = new StringBuilder(keyPrefix);
		
		while (!stack.empty()) {
			StackFrame<V> f = stack.peek();
			Node<V> n = f.node;
			
			Node<V> next = null;
			switch (++f.pos) {
			case StackFrame.LEFT:
				next = n.left;
				break;
			case StackFrame.MID:
				next = n.mid;
				
				if (n.value != null) {
					visitor.visit(keyBuilder.toString() + n.c, n.value);
				}
				
				if (next != null) {
					keyBuilder.append(n.c);
				}
				
				break;
			case StackFrame.RIGHT:
				next = n.right;
				break;
			}
			
			if (f.pos > StackFrame.RIGHT) {
				f = stack.pop();
				if (f.node.isMiddleChild()) {
					keyBuilder.deleteCharAt(keyBuilder.length() - 1);
				}
			}
			else if (next != null) {
				stack.push(new StackFrame<V>(next));
			}
		}
		
	}
	
	
	private static interface Visitor<V> {
		public void visit(String key, V value);
	}
	
	private static class StackFrame<V> {
		public final static int LEFT = 0;
		public final static int MID = 1;
		public final static int RIGHT = 2;
		
		private final Node<V> node;
		private int pos;
		
		public StackFrame(Node<V> node) {
			this.node = node;
			this.pos = -1;
		}
	}
	
	
	private static class Node<V> {
	
		private final char c;
		private V value;
		
		private final Node<V> parent; // parent is to useful when deleting
		private Node<V> left, right, mid;
		
		private Node(char c, Node<V> parent) {
			this.c = c;
			this.parent = parent;
		}

		public boolean isEmpty() {
			return value == null && left == null && right == null && mid == null;
		}
		
		private void removeChild(Node<V> child) {
			if (left == child) {
				left = null;
			}
			else if (right == child) {
				right = null;
			}
			else if (mid == child) {
				mid = null;
			}
		}
		
		public boolean isMiddleChild() {
			return parent != null && parent.mid == this;
		}
		
		@Override
		public String toString() {
			return "" + c + ":" + (value == null ? "." : value.toString());
		}
	}



	
	
	
}
